"""
Author: runyu.ma

parser.py
"""

from collections import namedtuple
from typing import List
# 假设 Token = namedtuple('Token', ['type','value','ln','col'])
Token = namedtuple('Token', ['type','value','ln','col'])

# AST 节点
class ASTNode:
    def __init__(self):
        self.parent: ASTNode = None
    
    def copy(self): return ASTNode()

class Number(ASTNode):
    def __init__(self, v: str): 
        super().__init__()
        self.value = v

    def __repr__(self): return f"Number({self.value})"
    def copy(self): return Number(str(self.value))

class String(ASTNode):
    def __init__(self, v: str, s: bool=False, p: bool=False):
        super().__init__()
        self.strong = s
        self.pystring = p
        self.value = v

    def __repr__(self): return f"String({self.value}, s={self.strong}, p={self.pystring})"
    def copy(self): return String(str(self.value), self.strong, self.pystring)

class Identifier(ASTNode):
    def __init__(self, n: str):
        super().__init__()
        self.name = n

    def __repr__(self): return f"Ident({self.name})"
    def copy(self): return Identifier(str(self.name))

class UnaryOp(ASTNode):
    def __init__(self, op: str, x: ASTNode):
        super().__init__()
        self.op, self.x = op, x
        if op == 'PTR' and x is None:
            self.x = String("Blank")
        self.x.parent = self

    def __repr__(self): return f"UnaryOp({self.op!r}, {self.x})"
    def copy(self): return UnaryOp(str(self.op), self.x.copy())

class BinaryOp(ASTNode):
    def __init__(self, op: str, l: ASTNode, r: ASTNode):
        super().__init__()
        self.op, self.l, self.r = op, l, r
        self.l.parent = self
        self.r.parent = self

    def __repr__(self): return f"BinaryOp({self.op!r}, {self.l}, {self.r})"
    def copy(self): return BinaryOp(str(self.op), self.l.copy(), self.r.copy())

class BlockOp(ASTNode):
    def __init__(self, oplist: List[ASTNode]):
        super().__init__()
        self.oplist = oplist
        for ch in self.oplist:
            ch.parent = self

    def __repr__(self): return f"BlockOp({self.oplist})"
    def copy(self): return BlockOp([e.copy() for e in self.oplist])

class ParserError(Exception): pass

class PrattParser:
    def __init__(self, tokens):
        # 在末尾加一个 EOF 方便终止
        self.tokens = tokens + [Token('EOF','',None,None)]
        self.pos    = -1

        # 三个 parselet 表
        self.atom_parselets   = {}  # 纯原子，不再调用 parse_expression
        self.prefix_parselets = {}  # 真正前缀运算符，后面要接右侧 expr
        self.infix_parselets  = {}  # 中缀运算符
        self.precedences      = {}  # 中缀的 binding power

    def current(self):
        return self.tokens[self.pos]

    def advance(self):
        self.pos += 1
        t = self.tokens[self.pos]
        return t
    
    def peek_nxt(self):
        return self.tokens[self.pos+1]

    def expect(self, ttype):
        t = self.current()
        #print("E", t)
        if t.type != ttype:
            raise ParserError(f"Expected {ttype}, got {t.type}")
        return self.advance()
    
    def expect_nxt(self, ttype):
        t = self.advance()
        #print("E", t)
        if t.type != ttype:
            raise ParserError(f"Expected {ttype}, got {t.type}")
        return t

    def parse(self):
        left_vec = [self.parse_expression()]
        while self.current().type != 'EOF' and self.peek_nxt().type != 'EOF':
            ast = self.parse_expression()
            if ast:
                left_vec.append(ast)
        ret = BlockOp(left_vec)
        return ret

    def parse_expression(self, rbp=0):
        # 1) 读取一个 token，先看是不是 atom
        tok = self.advance()
        #print("head", tok)
        if tok.type in ['SEMICOLON', 'RPAREN', 'END']:
            return None
        if tok.type in self.atom_parselets:
            left = self.atom_parselets[tok.type](tok)
            #if tok.type != 'LPAREN':
            #    self.advance()
        # 2) 再看是不是 prefix（真正需要消费右侧 expr 的）
        elif tok.type in self.prefix_parselets:
            left = self.prefix_parselets[tok.type](tok)
            if isinstance(left, UnaryOp) and left.op == 'PTR' and isinstance(left.x, String):
                print(f'PTR standalone, treat as type {self.tokens[self.pos-1]} {self.current()} {self.peek_nxt()}')
                #下一个END被消耗掉了，ptr作为type只有一种用法就是在 begin xxxx ptr end 中，后面一定是end
                self.pos = self.pos - 1
        #else:
        #    raise ParserError(f"Unexpected token {tok.type!r} at pos {self.pos}")
        #print('-------')
        #print(tok, left, self.current())
        # 3) 进入中缀循环
        while True:
            look = self.peek_nxt()
            #print("s", look)
            if look.type == 'EOF':
                break
            if look.type == 'SEMICOLON':
                break
            if look.type == 'RPAREN':
                break
            lbp = self.precedences.get(look.type+look.value, 0)
            #print("a", left, look, lbp, rbp)
            if lbp <= rbp:
                break
            self.advance()
            left = self.infix_parselets[look.type+look.value](left, look)
            #print("b", left, look)
        return left

    # 注册「原子」
    def register_atom(self, ttype, func):
        # func: Token -> ASTNode
        self.atom_parselets[ttype] = func

    # 注册「前缀」
    def register_prefix(self, ttype, precedence, func):
        # func: (Token, ASTNode) -> ASTNode
        # precedence 控制右侧 parse_expression(precedence)
        def _parselet(tok):
            right = self.parse_expression(precedence)
            return func(tok, right)
        self.prefix_parselets[ttype] = _parselet

    # 注册「中缀」
    def register_infix(self, ttype, tval, precedence, func, assoc='left'):
        # func: (leftAST, Token, rightAST) -> ASTNode
        def _parselet(left, tok):
            # 左结合：rbp = prec；右结合：rbp = prec-1
            next_rbp = precedence - 1 if assoc == 'right' else precedence
            right = self.parse_expression(next_rbp)
            return func(left, tok, right)
        self.infix_parselets[ttype+tval] = _parselet
        self.precedences[ttype+tval] = precedence


# ---- 举例：构建一个支持自定义运算符优先级的 parser ----
def build_default_parser(tokens):
    p = PrattParser(tokens)

    # 原子：数字、标识符、括号分组
    p.register_atom('NUMBER', lambda t: Number(t.value))
    p.register_atom('STRING', lambda t: String(t.value))
    p.register_atom('SSTRING',lambda t: String(t.value, s=True))
    p.register_atom('PSTRING',lambda t: String(t.value, p=True))
    p.register_atom('ID',     lambda t: Identifier(t.value))
    def parse_group(tok):
        expr = p.parse_expression()
        #print("P", expr)
        if expr:
            p.expect_nxt('RPAREN')
        else:
            p.expect('RPAREN')
        #print(p.current())
        return expr
    p.register_atom('LPAREN', parse_group)

    def parse_block(tok):
        expr_list = []
        while p.peek_nxt().type != 'END' and p.peek_nxt().type != 'EOF':
            expr = p.parse_expression()
            #print("B", expr)
            #print(p.current(), p.peek_nxt())
            if expr:
                expr_list.append(expr)
        #消耗掉最后的end，保证嵌套begin时，外层的begin不会被这次的end结束
        p.advance()
        return BlockOp(expr_list)
    p.register_atom('BEGIN', parse_block)

    # 前缀运算符：-x, +x
    p.register_prefix('MINUS',  80, lambda t, r: UnaryOp('-', r))
    p.register_prefix('PLUS',   80, lambda t, r: UnaryOp('+', r))
    p.register_prefix('PTR',    90, lambda t, r: UnaryOp('PTR', r))
    #p.register_prefix('DEREF',  90, lambda t, r: UnaryOp('DEREF', r))
    p.register_prefix('COLON',  15, lambda t, r: UnaryOp(':', r))
    p.register_prefix('RETURN', 90, lambda t, r: UnaryOp('RETURN', r))

    # 中缀运算符（可自定义优先级和结合性）
    ops = [
      ('OP',    '==',   10, 'left'),
      ('OP',    '=',    20, 'left'),
      ('OP',    '+',    40, 'left'),
      ('OP',    '-',    40, 'left'),
      ('OP',    '*',    50, 'left'),
      ('OP',    '/',    50, 'left'),
      ('OP',    '&',    30, 'left'),
      ('OP',    '|',    30, 'left'),
      ('OP',    '<<',   30, 'left'),
      ('OP',    '>>',   30, 'left'),
      ('OP',    '<',    10, 'left'),
      ('OP',    '>',    10, 'left'),
      ('OP',    '$',    80, 'left'),
      ('CAST',  '@',    80, 'left'),
      ('OP',    '!',    21,  'left'),
      ('COMMA', ',',    9,  'left'),
      # … 你可以再插入 ('^', 60, 'right') 之类
    ]
    for tok, val, prec, assoc in ops:
        p.register_infix(
          tok, val, prec,
          lambda l,t,r: BinaryOp(t.value, l, r),
          assoc=assoc
        )

    return p
